# -*- coding: utf-8 -*-

import os
import xml.etree.cElementTree as ET
import DbUtil as db_util


def pubmed_xml_parser(path, batch_size):
    dir_name = path.split('\\')[0]
    print(path)
    etree = ET.parse(path)
    root = etree.getroot()
    article_list = []
    journal_node_list, journal_article_rel_list = [], []
    medline_journal_node_info_list, medline_journal_rel_list, medline_article_rel_list = [], [], []
    mesh_words_node_list, mesh_words_article_rel_list = tuple(), tuple()
    chemical_words_node_list, chemical_words_article_rel_list = tuple(), tuple()
    pub_type_node_list, pub_type_article_rel_list = tuple(), tuple()
    reference_node_list, reference_article_rel_list = tuple(), tuple()
    author_node_list, author_article_re_list = tuple(), tuple()
    pmid_set = []
    count = 0
    author_auto_id = db_util.find_max_id('author')
    chemical_auto_id = db_util.find_max_id('chemical')
    mesh_word_auto_id = db_util.find_max_id('mesh_word')
    id_dict = {'author_auto_id': author_auto_id, 'chemical_auto_id': chemical_auto_id,
               'mesh_word_auto_id': mesh_word_auto_id}

    for article in root.iter('PubmedArticle'):
        pmid = article.find('MedlineCitation').find('PMID').text
        if pmid in pmid_set:
            continue
        pmid_set.append(pmid)

        article_data_list = parse_article(article)
        article_list.append(tuple(article_data_list))

        journal_node_data_list, journal_rel_data_list = parse_journal(article=article)
        if len(journal_node_data_list) > 0:
            journal_node_list.append(tuple(journal_node_data_list))
        if len(journal_rel_data_list) > 0:
            journal_article_rel_list.append(tuple(journal_rel_data_list))

        medline_journal_info_data_list, medline_journal_data_rel_list, medline_article_data_rel_list = parse_medline_journal_info(
            article=article)
        medline_journal_node_info_list.append(tuple(medline_journal_info_data_list))
        medline_journal_rel_list.append(tuple(medline_journal_data_rel_list))
        medline_article_rel_list.append(tuple(medline_article_data_rel_list))

        mesh_words_data_list, mesh_words_article_data_rel_list, mesh_word_max_id = parse_mesh_heading(article=article,
                                                                                                      current_id=
                                                                                                      id_dict[
                                                                                                          'mesh_word_auto_id'])
        id_dict['mesh_word_auto_id'] = mesh_word_max_id
        mesh_words_node_list = (*mesh_words_node_list, *mesh_words_data_list)
        mesh_words_article_rel_list = (*mesh_words_article_rel_list, *mesh_words_article_data_rel_list)

        chemical_words_data_list, chemical_words_article_data_rel_list, chemical_max_id = parse_chemical(
            article=article, current_id=id_dict['chemical_auto_id'])
        id_dict['chemical_auto_id'] = chemical_max_id
        chemical_words_node_list = (*chemical_words_node_list, *chemical_words_data_list)
        chemical_words_article_rel_list = (*chemical_words_article_rel_list, *chemical_words_article_data_rel_list)

        pub_type_data_list, pub_type_article_data_rel_list = parse_publication_type(article=article)
        pub_type_node_list = (*pub_type_node_list, *pub_type_data_list)
        pub_type_article_rel_list = (*pub_type_article_rel_list, *pub_type_article_data_rel_list)

        reference_data_list, reference_article_data_rel_list = parse_reference(article=article)
        reference_node_list = (*reference_node_list, *reference_data_list)
        reference_article_rel_list = (*reference_article_rel_list, *reference_article_data_rel_list)

        author_data_list, author_article_data_re_list, max_id = parse_author(article=article,
                                                                             current_id=id_dict['author_auto_id'])
        id_dict['author_auto_id'] = max_id
        author_node_list = (*author_node_list, *author_data_list)
        author_article_re_list = (*author_article_re_list, *author_article_data_re_list)

        count += 1
        if count >= batch_size:
            db_util.insert_article_node(article_list)

            db_util.insert_journal_node(journal_node_list)
            db_util.insert_article2journal_relation(journal_article_rel_list)

            db_util.insert_medline_journal_info_node(medline_journal_node_info_list)
            db_util.insert_medline2article_relation(medline_article_rel_list)
            db_util.insert_medline2journal_relation(medline_journal_rel_list)

            db_util.insert_mesh_word_node(mesh_words_node_list)
            db_util.insert_mesh_word2article_relation(mesh_words_article_rel_list)

            db_util.insert_chemical_node(chemical_words_node_list)
            db_util.insert_article2chemical_relation(chemical_words_article_rel_list)

            db_util.insert_publication_type_node(pub_type_node_list)
            db_util.insert_article2publication_type_relation(pub_type_article_rel_list)

            db_util.insert_reference_node(reference_node_list)
            db_util.insert_article2reference_relation(reference_article_rel_list)

            db_util.insert_author_node(author_node_list)
            db_util.insert_article2author_relation(author_article_re_list)

            article_list = []
            journal_node_list, journal_article_rel_list = [], []
            medline_journal_node_info_list, medline_journal_rel_list, medline_article_rel_list = [], [], []
            mesh_words_node_list, mesh_words_article_rel_list = tuple(), tuple()
            chemical_words_node_list, chemical_words_article_rel_list = tuple(), tuple()
            pub_type_node_list, pub_type_article_rel_list = tuple(), tuple()
            reference_node_list, reference_article_rel_list = tuple(), tuple()
            author_node_list, author_article_re_list = tuple(), tuple()

            count = 0
        print('pmid: ' + pmid + '\t done!')

    if len(article_list) > 0:
        db_util.insert_article_node(article_list)

        db_util.insert_journal_node(journal_node_list)
        db_util.insert_article2journal_relation(journal_article_rel_list)

        db_util.insert_medline_journal_info_node(medline_journal_node_info_list)
        db_util.insert_medline2article_relation(medline_article_rel_list)
        db_util.insert_medline2journal_relation(medline_journal_rel_list)

        db_util.insert_mesh_word_node(mesh_words_node_list)
        db_util.insert_mesh_word2article_relation(mesh_words_article_rel_list)

        db_util.insert_chemical_node(chemical_words_node_list)
        db_util.insert_article2chemical_relation(chemical_words_article_rel_list)

        db_util.insert_publication_type_node(pub_type_node_list)
        db_util.insert_article2publication_type_relation(pub_type_article_rel_list)

        db_util.insert_reference_node(reference_node_list)
        db_util.insert_article2reference_relation(reference_article_rel_list)

        db_util.insert_author_node(author_node_list)
        db_util.insert_article2author_relation(author_article_re_list)


def parse_author(article, current_id):
    '''
    解析作者，作者与文章关系数据
    :param current_id:
    :param article:
    :return:
    '''
    medline_citation_obj = article.find('MedlineCitation')
    pmid = medline_citation_obj.find('PMID').text
    author_data_list = []
    author_article_re_list = []
    max_id = current_id + 1 if current_id is not None else 1
    if medline_citation_obj.find('Article').find('AuthorList') is not None:
        authors = medline_citation_obj.find('Article').find('AuthorList').findall('Author')
        # auto_id = db_util.find_max_id()
        # max_id = auto_id if auto_id is not None else 0
        count = 1
        for author in authors:
            author_list = []
            author_article_list = []
            full_name, last_name, fore_name, initials = None, None, None, None
            if author.find('LastName') is not None:
                last_name = author.find('LastName').text
            if author.find('ForeName') is not None:
                fore_name = author.find('ForeName').text
            if author.find('Initials') is not None:
                initials = author.find('Initials').text

            author_id = 'author_' + str(max_id)
            author_list.append(author_id)
            if fore_name is not None and last_name is not None:
                full_name = fore_name + "." + last_name
            elif fore_name is not None:
                full_name = fore_name
            elif last_name is not None:
                full_name = last_name

            if count < len(authors):
                max_id += 1
                count += 1
            author_list.append(full_name)
            author_list.append(initials)
            author_data_list.append(tuple(author_list))

            author_article_list.append(pmid)
            author_article_list.append(author_id)
            author_article_re_list.append(tuple(author_article_list))
    return author_data_list, author_article_re_list, max_id


def parse_reference(article):
    '''
    解析参考文献，参考文献与文章关系数据
    :param article:
    :return:
    '''
    reference_data_list = []
    reference_article_rel_list = []
    medline_citation_obj = article.find('MedlineCitation')
    pmid = medline_citation_obj.find('PMID').text

    reference_obj = article.find('PubmedData').find('ReferenceList')
    if reference_obj != None:
        citation, reference_article_pmid = None, None

        for reference in reference_obj.iter('Reference'):
            citation = reference.find('Citation').text
            # <Citation>World Psychiatry. 2015 Feb;14(1):36-42</Citation>
            row = citation.split('.')
            for id_obj in reference.find('ArticleIdList').findall('ArticleId'):
                id_type = id_obj.attrib.get('IdType')
                if id_type == 'pubmed':
                    reference_article_pmid = id_obj.text
            reference_list, reference_article_list = [], []
            reference_list.append(reference_article_pmid)
            reference_list.append(row[0])  # article title
            reference_list.append(citation)
            reference_article_list.append(pmid)
            reference_article_list.append(reference_article_pmid)
            reference_article_list.append(citation)
            reference_data_list.append(tuple(reference_list))
            reference_article_rel_list.append(tuple(reference_article_list))

    return reference_data_list, reference_article_rel_list


def parse_publication_type(article):
    '''
    解析出版类型，出版类型与文章关系数据
    :param article:
    :return:
    '''
    medline_citation_obj = article.find('MedlineCitation')
    pmid = medline_citation_obj.find('PMID').text
    pub_type_article_rel_list = []
    pub_type_data_list = []
    for x in medline_citation_obj.find('Article').find('PublicationTypeList').findall('PublicationType'):
        publication_type_id = x.attrib.get('UI')
        publication_type = x.text
        # pub_type_text = publication_type_id + '/' + publication_type
        pub_type_list = []
        pub_type_article_list = []
        pub_type_list.append(publication_type_id)
        pub_type_list.append(publication_type)
        pub_type_article_list.append(pmid)
        pub_type_article_list.append(publication_type_id)
        pub_type_data_list.append(tuple(pub_type_list))
        pub_type_article_rel_list.append(tuple(pub_type_article_list))

    return pub_type_data_list, pub_type_article_rel_list


def parse_chemical(article, current_id):
    '''
    解析化学品、化学品与文章的关系数据
    :param current_id:
    :param article:
    :return:
    '''

    chemical_words_data_list = []
    chemical_words_article_rel_list = []
    medline_citation_obj = article.find('MedlineCitation')
    pmid = medline_citation_obj.find('PMID').text

    max_id = current_id + 1 if current_id is not None else 1
    count = 1
    if medline_citation_obj.find('ChemicalList') is not None:
        chemicals = medline_citation_obj.find('ChemicalList').findall('Chemical')
        for chemical in chemicals:
            chemical_id, name_of_substance, registry_number = None, None, None
            chemical_words_list = []
            chemical_words_article_list = []
            if chemical.find('NameOfSubstance') is not None:
                chemical_id = chemical.find('NameOfSubstance').attrib.get('UI')
                name_of_substance = chemical.find('NameOfSubstance').text
            if chemical.find('RegistryNumber') is not None:
                registry_number = chemical.find('RegistryNumber').text

            chemical_id = 'ch-' + str(max_id)
            chemical_words_list.append(chemical_id)
            chemical_words_list.append(name_of_substance)
            chemical_words_list.append(registry_number)
            chemical_words_data_list.append(tuple(chemical_words_list))

            chemical_words_article_list.append(pmid)
            chemical_words_article_list.append(chemical_id)
            chemical_words_article_rel_list.append(tuple(chemical_words_article_list))
            if count < len(chemicals):
                max_id += 1
                count += 1

    return chemical_words_data_list, chemical_words_article_rel_list, max_id


def parse_mesh_heading(article, current_id):
    '''
    解析主题词、主题词与文章关系数据
    :param current_id:
    :param article:
    :return:
    '''

    mesh_words_data_list = []
    mesh_words_article_rel_list = []
    medline_citation_obj = article.find('MedlineCitation')
    pmid = medline_citation_obj.find('PMID').text

    max_id = current_id + 1 if current_id is not None else 1
    count = 1
    mesh_words_objects = medline_citation_obj.find('MeshHeadingList').findall('MeshHeading')
    for mesh_heading in mesh_words_objects:
        mesh_heading_id, descriptor_name, qualifier_name = None, None, None
        mesh_words_list = []
        mesh_words_article_list = []
        if mesh_heading.find('DescriptorName') is not None:
            mesh_heading_id = mesh_heading.find('DescriptorName').attrib.get('UI')
            descriptor_name = mesh_heading.find('DescriptorName').text
        if mesh_heading.find('QualifierName') is not None:
            qualifier_name_id = mesh_heading.find('QualifierName').attrib.get('UI')
            qualifier_name = mesh_heading.find('QualifierName').text

        mesh_heading_id = 'mw-' + str(max_id)
        mesh_words_list.append(mesh_heading_id)
        mesh_words_list.append(descriptor_name)
        mesh_words_list.append(qualifier_name)
        mesh_words_data_list.append(tuple(mesh_words_list))

        mesh_words_article_list.append(pmid)
        mesh_words_article_list.append(mesh_heading_id)
        mesh_words_article_rel_list.append(tuple(mesh_words_article_list))
        if count < len(mesh_words_objects):
            max_id += 1
            count += 1

    return mesh_words_data_list, mesh_words_article_rel_list, max_id


def parse_medline_journal_info(article):
    '''
    解析medline对象和medline与journal的关系数据、medline与article的关系数据
    '''

    medline_journal_info_data_list = []
    medline_journal_rel_list = []
    medline_article_rel_list = []

    medline_citation_obj = article.find('MedlineCitation')
    pmid = medline_citation_obj.find('PMID').text

    medline_journal_info_obj = medline_citation_obj.find('MedlineJournalInfo')
    country, medline_ta, nlm_unique_id, issn_linking = None, None, None, None
    if medline_journal_info_obj.find('Country') is not None:
        country = medline_journal_info_obj.find('Country').text
    if medline_journal_info_obj.find('MedlineTA') is not None:
        medline_ta = medline_journal_info_obj.find('MedlineTA').text
    if medline_journal_info_obj.find('NlmUniqueID') is not None:
        nlm_unique_id = medline_journal_info_obj.find('NlmUniqueID').text
    if medline_journal_info_obj.find('ISSNLinking') is not None:
        issn_linking = medline_journal_info_obj.find('ISSNLinking').text

    # journal = medline_citation_obj.find('Article').find('Journal')
    # isbn = None
    # if journal.find('ISSN') is not None:
    #     isbn = journal.find('ISSN').text

    medline_journal_info_data_list.append(nlm_unique_id)
    medline_journal_info_data_list.append(nlm_unique_id)
    medline_journal_info_data_list.append(country)
    medline_journal_info_data_list.append(medline_ta)
    medline_journal_info_data_list.append(issn_linking)

    medline_journal_rel_list.append(nlm_unique_id)
    medline_journal_rel_list.append(issn_linking)

    medline_article_rel_list.append(pmid)
    medline_article_rel_list.append(nlm_unique_id)
    return medline_journal_info_data_list, medline_journal_rel_list, medline_article_rel_list


def parse_journal(article):
    '''
    解析杂志journal对象
    :param article: dom中的journal对象
    :return:journal实体数据和article与journal的关系数据
    '''
    volume, issue = None, None
    journal_node_data_list = []
    journal_rel_data_list = []
    pmid = article.find('MedlineCitation').find('PMID').text

    journal = article.find('MedlineCitation').find('Article').find('Journal')
    if journal.find('ISSN') is not None:
        issn = journal.find('ISSN').text
        journal_rel_data_list.append(pmid)
        journal_rel_data_list.append(issn)

        volume_obj = journal.find('JournalIssue').find("Volume")
        if volume_obj is not None:
            volume = volume_obj.text
        issue_obj = journal.find('JournalIssue').find("Issue")
        if issue_obj is not None:
            issue = issue_obj.text

        volume_issue = None
        if volume is not None and issue is not None:
            volume_issue = volume + '-' + issue
        elif volume is not None:
            volume_issue = volume
        elif issue is not None:
            volume_issue = issue

        journal_rel_data_list.append(volume_issue)

        iso_abbreviation = journal.find('ISOAbbreviation').text
        journal_title = journal.find('Title').text

        pub_date_obj = journal.find('JournalIssue').find('PubDate')
        date = None
        if pub_date_obj.find('MedlineDate') is not None:
            date = pub_date_obj.find('MedlineDate').text
        else:
            year_obj = pub_date_obj.find('Year')
            month_obj = pub_date_obj.find('Month')
            day_obj = pub_date_obj.find('Day')
            if year_obj is not None:
                date = year_obj.text
            if month_obj is not None:
                date = date + '-' + month_obj.text
            if day_obj is not None:
                date = date + '-' + day_obj.text

        journal_node_data_list.append(issn)
        journal_node_data_list.append(issn)
        journal_node_data_list.append(journal_title)
        journal_node_data_list.append(iso_abbreviation)
        journal_node_data_list.append(volume)
        journal_node_data_list.append(issue)
        journal_node_data_list.append(date)

    return journal_node_data_list, journal_rel_data_list


def parse_article(article):
    '''
    解析文章article对象
    :param article: dom中的一篇文章对象
    :return:
    '''
    article_data_list = []
    pmid = article.find('MedlineCitation').find('PMID').text
    article_data_list.append(pmid)

    # ArticleIdList
    id_map = {}
    for article_id in article.find('PubmedData').find('ArticleIdList').findall('ArticleId'):
        id_type = article_id.attrib.get('IdType')
        aid = article_id.text
        id_map[id_type] = aid
        # print(id_type + ":" + aid)

    doi = id_map['doi'] if 'doi' in id_map is not None else None
    pii = id_map['pii'] if 'pii' in id_map is not None else None
    pmc = id_map['pmc'] if 'pmc' in id_map is not None else None
    article_data_list.append(doi)
    article_data_list.append(pii)
    article_data_list.append(pmc)

    article_object = article.find('MedlineCitation').find('Article')
    article_title = article_object.find('ArticleTitle').text
    language = article_object.find('Language').text
    medline_pgn = article_object.find('Pagination').find('MedlinePgn').text
    article_data_list.append(remove_useless(article_title))
    article_data_list.append(medline_pgn)
    # PublicationStatus
    publication_status = article.find('PubmedData').find('PublicationStatus').text
    article_data_list.append(publication_status)
    article_data_list.append(language)
    abstract, vernacular_title, citation_subset = None, None, None
    if article_object.find('Abstract') is not None:
        abstract = article_object.find('Abstract').find('AbstractText').text

    if article_object.find('VernacularTitle') is not None:
        vernacular_title = article_object.find('VernacularTitle').text

    if article.find('MedlineCitation').find('CitationSubset') is not None:
        citation_subset = article.find('MedlineCitation').find('CitationSubset').text
    article_data_list.append(remove_useless(vernacular_title))
    article_data_list.append(remove_useless(abstract))
    article_data_list.append(citation_subset)

    return article_data_list


def remove_useless(text):
    if text is not None:
        text = text.replace('"', '')
    return text


def get_files_path(base_path_local):
    for base_path, folders, files in os.walk('D:\jobs\data'):
        file_list = [os.path.join(base_path_local, file) for file in files if file.endswith('.xml')]

    print('共需解析文件：{0}'.format(len(file_list)))
    return file_list


if __name__ == '__main__':
    batch_size = 2000
    base_path_local = "D:\jobs\data"
    files_list = get_files_path(base_path_local=base_path_local)
    if not files_list:
        print('全部解析完成')
    else:
        for path in files_list:
            pubmed_xml_parser(path, batch_size)

        print('all articles has parsed succeed!')
        # with ThreadPoolExecutor(max_workers=os.cpu_count()) as pool:
        #     pool.map(pubmed_xml_parser, files_list)
